package com.tuniu.agents.evaluators.common.impl;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class FavoritesEvaluatorStreamHandle {

    private static final String THINK_START_TAG = "<think>";
    private static final String THINK_END_TAG = "</think>";
    private final Map<String, StreamState> streamStates = new ConcurrentHashMap<>();

    private String conversationId;

    public FavoritesEvaluatorStreamHandle(String conversationId) {
        this.conversationId = conversationId;
    }

    public String handleMessageWithThink(String chunk) {

        StreamState state = streamStates.computeIfAbsent(conversationId, k -> new StreamState());
        state.buffer.append(chunk);

        StringBuilder result = new StringBuilder();
        StringBuilder currentBuffer = state.buffer;
        int i = 0;
        while (i < currentBuffer.length()) {
            if (!state.inThink) {
                int startThink = currentBuffer.indexOf(THINK_START_TAG, i);
                if (startThink == -1) {
                    result.append(currentBuffer.substring(i));
                    i = currentBuffer.length();
                } else {
                    result.append(currentBuffer.substring(i, startThink));
                    state.inThink = true;
                    i = startThink + THINK_START_TAG.length();
                }
            } else {
                int endThink = currentBuffer.indexOf(THINK_END_TAG, i);
                if (endThink == -1) {
                    i = currentBuffer.length();
                } else {
                    state.inThink = false;
                    i = endThink + THINK_END_TAG.length();
                }
            }
        }

        String remaining = currentBuffer.substring(i);
        currentBuffer.setLength(0);
        currentBuffer.append(remaining);

        return result.toString();
    }

    private static class StreamState {
        final StringBuilder buffer = new StringBuilder();
        boolean inThink = false;
    }

}
